#include "BYTETracker.h"
#include <fstream>

std::shared_ptr<BYTETracker> BYTETracker::tracker_ = nullptr;

BYTETracker::BYTETracker()
{
    // track_thresh = 0.5;
    // high_thresh = 0.6;
    // match_thresh = 0.8;

    // frame_id = 0;
    // max_time_lost = int(frame_rate / 30.0 * track_buffer);
    // cout << "Init ByteTrack!" << endl;
}

bool BYTETracker::init(double track_match_thresh, int track_max_time_lost)
{
    track_thresh = 0.1;
    high_thresh = 0.25;
    match_thresh = track_match_thresh;

    frame_id = 0;
    max_time_lost = track_max_time_lost; // int(frame_rate / 30.0 * track_buffer);
    std::cout << "Init ByteTrack! track_max_time_lost!:" << track_max_time_lost
              << ",match_thresh:" << match_thresh << std::endl;
    return true;
}

int BYTETracker::update(const std::vector<perception::camera::ObjectDetectInfoPtr>& objects,
    std::vector<perception::camera::ObjectTrackInfoPtr>& tracked_objects, const cv::Mat&, int64_t)
{

    ////////////////// Step 1: Get detections //////////////////
    this->frame_id++;
    std::vector<STrack> activated_stracks;
    std::vector<STrack> refind_stracks;
    std::vector<STrack> removed_stracks;
    std::vector<STrack> lost_stracks;
    std::vector<STrack> detections;
    std::vector<STrack> detections_low;

    std::vector<STrack> detections_cp;
    std::vector<STrack> tracked_stracks_swap;
    std::vector<STrack> resa, resb;
    // std::vector<STrack> output_stracks;

    std::vector<STrack*> unconfirmed;
    std::vector<STrack*> tracked_stracks;
    std::vector<STrack*> strack_pool;
    std::vector<STrack*> r_tracked_stracks;
    // static int _trackid_ = 0;
    // for (size_t i = 0; i < objects.size(); i++) {
    //     perception::camera::ObjectTrackInfoPtr oti = std::make_shared<perception::camera::ObjectTrackInfo>();
    //     oti->detect.type_id = objects[i]->type_id;
    //     oti->detect.type_id_confidence = objects[i]->type_id_confidence;
    //     oti->detect.bottom_uv = objects[i]->bottom_uv;
    //     oti->detect.pts8 = objects[i]->pts8;
    //     oti->detect.box.left_top = objects[i]->box.left_top;
    //     oti->detect.box.right_bottom = objects[i]->box.right_bottom;

    //     oti->tracker.track_id = _trackid_ ++;
    //     tracked_objects.push_back(oti);
    // }
    // return 0;

    if (objects.size() > 0) {
        for (size_t i = 0; i < objects.size(); i++) {
            std::vector<float> tlbr_;
            tlbr_.resize(4);
            tlbr_[0] = objects[i]->box.left_top.x;
            tlbr_[1] = objects[i]->box.left_top.y;
            tlbr_[2] = objects[i]->box.right_bottom.x;
            tlbr_[3] = objects[i]->box.right_bottom.y;

            float score = objects[i]->type_id_confidence;

            STrack strack(STrack::tlbr_to_tlwh(tlbr_), score, objects[i]->type_id);
            strack.pts8 = objects[i]->pts8;
            strack.object_feature = objects[i]->object_feature;
            strack.bottom_uv = objects[i]->bottom_uv;
            if (score >= track_thresh) {
                // cv::rectangle(img, cv::Point2d(strack.tlbr[0], strack.tlbr[1]),
                //                     cv::Point2d(strack.tlbr[2], strack.tlbr[3]), cv::Scalar{0, 0, 255});
                detections.push_back(strack);
            } else {
                // cv::rectangle(img, cv::Point2d(strack.tlbr[0], strack.tlbr[1]),
                //                     cv::Point2d(strack.tlbr[2], strack.tlbr[3]), cv::Scalar{0, 255, 0});
                detections_low.push_back(strack);
            }
        }
    }

    // Add newly detected tracklets to tracked_stracks
    for (size_t i = 0; i < this->tracked_stracks.size(); i++) {
        if (!this->tracked_stracks[i].is_activated) {
            unconfirmed.push_back(&this->tracked_stracks[i]);
            // cv::rectangle(img, cv::Point2d(this->tracked_stracks[i].tlbr[0], this->tracked_stracks[i].tlbr[1]),
            //                         cv::Point2d(this->tracked_stracks[i].tlbr[2], this->tracked_stracks[i].tlbr[3]),
            //                         cv::Scalar{255, 0, 0});
            // cv::putText(img, std::string("noact"), cv::Point2d(this->tracked_stracks[i].tlbr[0], this->tracked_stracks[i].tlbr[1]-5),
            //             cv::LINE_AA, 0.5, cv::Scalar{255, 0, 0} );

        } else {
            tracked_stracks.push_back(&this->tracked_stracks[i]);
            // cv::rectangle(img, cv::Point2d(this->tracked_stracks[i].tlbr[0], this->tracked_stracks[i].tlbr[1]),
            //                         cv::Point2d(this->tracked_stracks[i].tlbr[2], this->tracked_stracks[i].tlbr[3]),
            //                         cv::Scalar{128, 128, 0});
            // cv::putText(img, std::string("act"), cv::Point2d(this->tracked_stracks[i].tlbr[0], this->tracked_stracks[i].tlbr[1]-5),
            //             cv::LINE_AA, 0.5, cv::Scalar{128, 128, 0} );
        }
    }

    ////////////////// Step 2: First association, with IoU //////////////////
    strack_pool = joint_stracks(tracked_stracks, this->lost_stracks);
    STrack::multi_predict(strack_pool, this->kalman_filter);

    std::vector<std::vector<float>> dists;
    int dist_size = 0, dist_size_size = 0;
    dists = iou_distance(strack_pool, detections, dist_size, dist_size_size);
    // for(int isp = 0; isp < strack_pool.size(); isp ++){
    //     cv::Mat img_draw = img.clone();
    //     cv::rectangle(img_draw, cv::Point2d(strack_pool.at(isp)->tlbr[0], strack_pool.at(isp)->tlbr[1]),
    //                     cv::Point2d(strack_pool.at(isp)->tlbr[2], strack_pool.at(isp)->tlbr[3]),
    //                     cv::Scalar{0, 0, 255});
    //     for(int idet = 0; idet < detections.size(); ++idet){
    //         cv::rectangle(img_draw, cv::Point2d(detections.at(idet).tlbr[0], detections.at(idet).tlbr[1]),
    //                     cv::Point2d(detections.at(idet).tlbr[2], detections.at(idet).tlbr[3]),
    //                     cv::Scalar{128, 0, 128});
    //         cv::putText(img_draw, std::to_string(dists.at(isp).at(idet)), cv::Point2d(detections.at(idet).tlbr[0], detections.at(idet).tlbr[1]-5),
    //                     cv::LINE_AA, 0.5, cv::Scalar{128, 0, 128} );
    //     }
    //     cv::imwrite("./var/log/visible/" + std::to_string(ts_ms) + "_" + std::to_string(isp) + "_tk.jpg", img_draw);
    // }

    std::vector<std::vector<int>> matches;
    std::vector<int> u_track, u_detection;
    // std::cout << "linear_assignment:" << dist_size << ", " << dist_size_size << ", match_thresh:" << match_thresh << std::endl;
    linear_assignment(dists, dist_size, dist_size_size, match_thresh, matches, u_track, u_detection);
    // for(int itk = 0; itk < u_track.size(); ++itk){
    //     int tid = u_track.at(itk);
    //     cv::rectangle(img, cv::Point2d(strack_pool.at(itk)->tlbr[0], strack_pool.at(itk)->tlbr[1]),
    //                         cv::Point2d(strack_pool.at(itk)->tlbr[2], strack_pool.at(itk)->tlbr[3]),
    //                         cv::Scalar{128, 0, 128});
    //     cv::putText(img, std::to_string(tid), cv::Point2d(strack_pool.at(itk)->tlbr[0], strack_pool.at(itk)->tlbr[1]-5),
    //                         cv::LINE_AA, 0.5, cv::Scalar{128, 0, 128} );
    // }
    // for(int iud = 0; iud < u_detection.size(); ++iud){
    //     cv::rectangle(img, cv::Point2d(detections.at(iud).tlbr[0], detections.at(iud).tlbr[1]),
    //                         cv::Point2d(detections.at(iud).tlbr[2], detections.at(iud).tlbr[3]),
    //                         cv::Scalar{0, 0, 255});
    //     // cv::putText(img, std::to_string(dists.at(isp).at(idet)), cv::Point2d(detections.at(iud).tlbr[0], detections.at(iud).tlbr[1]-5),
    //     //                     cv::LINE_AA, 0.5, cv::Scalar{128, 0, 128} );
    // }
    // cv::imwrite("./var/log/visible/" + std::to_string(ts_ms) + "_" + std::to_string(this->frame_id) + "_unmatch.jpg", img);

    for (size_t i = 0; i < matches.size(); i++) {
        STrack* track = strack_pool[matches[i][0]];
        STrack* det = &detections[matches[i][1]];
        if (track->state == TrackState::TRACKED) {
            track->update(*det, this->frame_id);
            activated_stracks.push_back(*track);
            // cv::rectangle(img, cv::Point2d(track->tlbr[0], track->tlbr[1]),
            //             cv::Point2d(track->tlbr[2], track->tlbr[3]),
            //             cv::Scalar{128, 0, 128});
            // cv::putText(img, std::string("update_track"), cv::Point2d(track->tlbr[0], track->tlbr[1]-5),
            //             cv::LINE_AA, 0.5, cv::Scalar{128, 0, 128} );

        } else {
            track->re_activate(*det, this->frame_id, false);
            refind_stracks.push_back(*track);
            // cv::rectangle(img, cv::Point2d(track->tlbr[0], track->tlbr[1]),
            //             cv::Point2d(track->tlbr[2], track->tlbr[3]),
            //             cv::Scalar{0, 128, 128});
            // cv::putText(img, std::string("react"), cv::Point2d(track->tlbr[0], track->tlbr[1]-5),
            //             cv::LINE_AA, 0.5, cv::Scalar{0, 128, 128} );
        }
    }

    ////////////////// Step 3: Second association, using low score dets //////////////////
    for (size_t i = 0; i < u_detection.size(); i++) {
        detections_cp.push_back(detections[u_detection[i]]);
    }
    detections.clear();
    detections.assign(detections_low.begin(), detections_low.end());

    for (size_t i = 0; i < u_track.size(); i++) {
        if (strack_pool[u_track[i]]->state == TrackState::TRACKED) {
            r_tracked_stracks.push_back(strack_pool[u_track[i]]);
            // cv::rectangle(img, cv::Point2d(strack_pool[u_track[i]]->tlbr[0], strack_pool[u_track[i]]->tlbr[1]),
            //             cv::Point2d(strack_pool[u_track[i]]->tlbr[2], strack_pool[u_track[i]]->tlbr[3]),
            //             cv::Scalar{0, 128, 255});
            // cv::putText(img, std::string("lowtracked"),
            //             cv::Point2d(strack_pool[u_track[i]]->tlbr[0], strack_pool[u_track[i]]->tlbr[1]-5),
            //             cv::LINE_AA, 0.5, cv::Scalar{0, 128, 255} );
        }
    }

    dists.clear();
    dists = iou_distance(r_tracked_stracks, detections, dist_size, dist_size_size);

    matches.clear();
    u_track.clear();
    u_detection.clear();
    linear_assignment(dists, dist_size, dist_size_size, 0.5, matches, u_track, u_detection);

    for (size_t i = 0; i < matches.size(); i++) {
        STrack* track = r_tracked_stracks[matches[i][0]];
        STrack* det = &detections[matches[i][1]];
        if (track->state == TrackState::TRACKED) {
            track->update(*det, this->frame_id);
            activated_stracks.push_back(*track);
            // cv::rectangle(img, cv::Point2d(track->tlbr[0], track->tlbr[1]),
            //             cv::Point2d(track->tlbr[2], track->tlbr[3]),
            //             cv::Scalar{128, 128, 128});
            // cv::putText(img, std::string("lowtrackmatch"), cv::Point2d(track->tlbr[0], track->tlbr[1]-5),
            //             cv::LINE_AA, 0.5, cv::Scalar{128, 128, 128} );

        } else {
            track->re_activate(*det, this->frame_id, false);
            refind_stracks.push_back(*track);
            // cv::rectangle(img, cv::Point2d(track->tlbr[0], track->tlbr[1]),
            //             cv::Point2d(track->tlbr[2], track->tlbr[3]),
            //             cv::Scalar{255, 128, 128});
            // cv::putText(img, std::string("lowreactmatch"), cv::Point2d(track->tlbr[0], track->tlbr[1]-5),
            //             cv::LINE_AA, 0.5, cv::Scalar{255, 128, 128} );
        }
    }

    for (size_t i = 0; i < u_track.size(); i++) {
        STrack* track = r_tracked_stracks[u_track[i]];
        if (track->state != TrackState::LOST) {
            track->mark_lost();
            lost_stracks.push_back(*track);
            // cv::rectangle(img, cv::Point2d(track->tlbr[0], track->tlbr[1]),
            //             cv::Point2d(track->tlbr[2], track->tlbr[3]),
            //             cv::Scalar{128, 128, 255});
            // cv::putText(img, std::string("lost"), cv::Point2d(track->tlbr[0], track->tlbr[1]-5),
            //             cv::LINE_AA, 0.5, cv::Scalar{128, 128, 255} );
        }
    }

    // Deal with unconfirmed tracks, usually tracks with only one beginning frame
    detections.clear();
    detections.assign(detections_cp.begin(), detections_cp.end());

    dists.clear();
    dists = iou_distance(unconfirmed, detections, dist_size, dist_size_size);

    matches.clear();
    std::vector<int> u_unconfirmed;
    u_detection.clear();
    // linear_assignment(dists, dist_size, dist_size_size, 0.7, matches, u_unconfirmed, u_detection);
    linear_assignment(dists, dist_size, dist_size_size, match_thresh, matches, u_unconfirmed, u_detection);

    for (size_t i = 0; i < matches.size(); i++) {
        unconfirmed[matches[i][0]]->update(detections[matches[i][1]], this->frame_id);
        activated_stracks.push_back(*unconfirmed[matches[i][0]]);
        // cv::rectangle(img, cv::Point2d(unconfirmed[matches[i][0]]->tlbr[0], unconfirmed[matches[i][0]]->tlbr[1]),
        //                 cv::Point2d(unconfirmed[matches[i][0]]->tlbr[2], unconfirmed[matches[i][0]]->tlbr[3]),
        //                 cv::Scalar{128, 0, 255});
        // cv::putText(img, std::string("lost"), cv::Point2d(unconfirmed[matches[i][0]]->tlbr[0], unconfirmed[matches[i][0]]->tlbr[1]-5),
        //             cv::LINE_AA, 0.5, cv::Scalar{128, 0, 255} );
    }

    for (size_t i = 0; i < u_unconfirmed.size(); i++) {
        STrack* track = unconfirmed[u_unconfirmed[i]];
        track->mark_removed();
        removed_stracks.push_back(*track);
    }

    ////////////////// Step 4: Init new stracks //////////////////
    for (size_t i = 0; i < u_detection.size(); i++) {
        STrack* track = &detections[u_detection[i]];
        if (track->score < this->high_thresh)
            continue;
        track->activate(this->kalman_filter, this->frame_id);
        activated_stracks.push_back(*track);
    }

    ////////////////// Step 5: Update state //////////////////
    // std::cout << "this->max_time_lost:" << this->max_time_lost << std::endl;
    for (size_t i = 0; i < this->lost_stracks.size(); i++) {
        if (this->frame_id - this->lost_stracks[i].end_frame() > this->max_time_lost) {
            this->lost_stracks[i].mark_removed();
            removed_stracks.push_back(this->lost_stracks[i]);
        }
    }

    for (size_t i = 0; i < this->tracked_stracks.size(); i++) {
        if (this->tracked_stracks[i].state == TrackState::TRACKED) {
            tracked_stracks_swap.push_back(this->tracked_stracks[i]);
        }
    }
    this->tracked_stracks.clear();
    this->tracked_stracks.assign(tracked_stracks_swap.begin(), tracked_stracks_swap.end());

    this->tracked_stracks = joint_stracks(this->tracked_stracks, activated_stracks);
    this->tracked_stracks = joint_stracks(this->tracked_stracks, refind_stracks);

    // std::cout << activated_stracks.size() << std::endl;

    this->lost_stracks = sub_stracks(this->lost_stracks, this->tracked_stracks);
    for (size_t i = 0; i < lost_stracks.size(); i++) {
        this->lost_stracks.push_back(lost_stracks[i]);
    }

    this->lost_stracks = sub_stracks(this->lost_stracks, this->removed_stracks);
    for (size_t i = 0; i < removed_stracks.size(); i++) {
        this->removed_stracks.push_back(removed_stracks[i]);
    }

    remove_duplicate_stracks(resa, resb, this->tracked_stracks, this->lost_stracks);

    this->tracked_stracks.clear();
    this->tracked_stracks.assign(resa.begin(), resa.end());
    this->lost_stracks.clear();
    this->lost_stracks.assign(resb.begin(), resb.end());

    for (size_t i = 0; i < this->tracked_stracks.size(); i++) {
        perception::camera::ObjectTrackInfoPtr oti = std::make_shared<perception::camera::ObjectTrackInfo>();
        oti->detect.type_id = this->tracked_stracks[i].clsid;
        oti->detect.type_id_confidence = this->tracked_stracks[i].score;
        oti->detect.bottom_uv = this->tracked_stracks[i].bottom_uv;
        oti->detect.pts8 = this->tracked_stracks[i].pts8;
        oti->detect.object_feature = this->tracked_stracks[i].object_feature;
        oti->detect.box.left_top = { this->tracked_stracks[i].tlbr[0], this->tracked_stracks[i].tlbr[1] };
        oti->detect.box.right_bottom = { this->tracked_stracks[i].tlbr[2], this->tracked_stracks[i].tlbr[3] };

        if (this->tracked_stracks[i].is_activated) {
            oti->tracker.track_id = this->tracked_stracks[i].track_id;
            // output_stracks.push_back(this->tracked_stracks[i]);

        } else {
            oti->tracker.track_id = -this->tracked_stracks[i].track_id;
        }
        tracked_objects.push_back(oti);
    }

    // for (int i = 0; i < this->lost_stracks.size(); i++) {
    //     if (this->lost_stracks[i].is_activated) {
    //         // output_stracks.push_back(this->tracked_stracks[i]);
    //         perception::camera::ObjectTrackInfoPtr oti = std::make_shared<perception::camera::ObjectTrackInfo>();
    //         oti->detect.type_id = this->lost_stracks[i].clsid;
    //         oti->detect.type_id_confidence = this->lost_stracks[i].score;
    //         oti->detect.bottom_uv = this->lost_stracks[i].bottom_uv;
    //         oti->detect.pts8 = this->lost_stracks[i].pts8;
    //         oti->detect.box.left_top = { this->lost_stracks[i].tlbr[0], this->lost_stracks[i].tlbr[1] };
    //         oti->detect.box.right_bottom = { this->lost_stracks[i].tlbr[2], this->lost_stracks[i].tlbr[3] };
    //         oti->tracker.track_id = -this->lost_stracks[i].track_id;
    //         tracked_objects.push_back(oti);
    //     }
    // }
    return 0;
}
